import tensorflow as tf
from tensorflow.keras import activations, initializers, constraints
from tensorflow.keras import regularizers
from tensorflow.keras.layers import Layer


class CrossGate(Layer):
    def __init__(self, units: int, features,
                 activation=None,
                 out_activation=None,
                 use_bias=True,
                 out_use_bias=False,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 **kwargs):
        if 'input_shape' not in kwargs and 'input_dim' in kwargs:
            kwargs['input_shape'] = (kwargs.pop('input_dim'),)
        super(CrossGate, self).__init__(**kwargs)
        self.units = units
        self.activation = activations.get(activation)
        self.out_activation = activations.get(out_activation)
        self.use_bias = use_bias
        self.out_use_bias = out_use_bias
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.bias_regularizer = regularizers.get(bias_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)
        self.bias_constraint = constraints.get(bias_constraint)
        self.features = tf.constant(features, dtype='float32')
        self.kernel = None
        self.bias = {}

    def build(self, input_shape):
        assert input_shape[1] == self.features.shape[0]
        feature_dim = self.features.shape[1]
        item_dim = input_shape[2]
        self.kernel = {'W_f': self.add_weight(shape=(feature_dim, item_dim),
                                              initializer=self.kernel_initializer,
                                              name='kernel_W_f',
                                              regularizer=self.kernel_regularizer,
                                              constraint=self.kernel_constraint),

                       'W_i2f': self.add_weight(shape=(item_dim, item_dim),
                                                initializer=self.kernel_initializer,
                                                name='kernel_W_i2f',
                                                regularizer=self.kernel_regularizer,
                                                constraint=self.kernel_constraint),
                       'W_f2i': self.add_weight(shape=(item_dim, item_dim),
                                                initializer=self.kernel_initializer,
                                                name='kernel_W_f2i',
                                                regularizer=self.kernel_regularizer,
                                                constraint=self.kernel_constraint),
                       'W_out': self.add_weight(shape=(item_dim * 2, self.units),
                                                initializer=self.kernel_initializer,
                                                name='kernel_W_out',
                                                regularizer=self.kernel_regularizer,
                                                constraint=self.kernel_constraint),
                       }
        if self.use_bias is True:
            self.bias['b_i2f'] = self.add_weight(shape=(item_dim,),
                                                 initializer=self.bias_initializer,
                                                 name='bias_i2f',
                                                 regularizer=self.bias_regularizer,
                                                 constraint=self.bias_constraint)
            self.bias['b_f2i'] = self.add_weight(shape=(item_dim,),
                                                 initializer=self.bias_initializer,
                                                 name='bias_f2i',
                                                 regularizer=self.bias_regularizer,
                                                 constraint=self.bias_constraint)
        if self.out_use_bias is True:
            self.bias['b_out'] = self.add_weight(shape=(self.units,),
                                                 initializer=self.bias_initializer,
                                                 name='bias_out',
                                                 regularizer=self.bias_regularizer,
                                                 constraint=self.bias_constraint)
        if len(self.bias) == 0:
            self.bias = None

    def call(self, inputs, **kwargs):
        item_embeddings = inputs[0]
        features = tf.matmul(self.features, self.kernel['W_f'])

        gate_f = tf.matmul(item_embeddings, self.kernel['W_i2f'])
        if self.use_bias:
            gate_f += self.bias['b_i2f']
        if self.activation is not None:
            gate_f = self.activation(gate_f)

        gate_i = tf.matmul(item_embeddings, self.kernel['W_f2i'])
        if self.use_bias:
            gate_i += self.bias['b_f2i']
        if self.activation is not None:
            gate_i = self.activation(gate_i)

        gated_features = tf.multiply(features, gate_f)
        gated_item = tf.multiply(item_embeddings, gate_i)

        outs = tf.concat([gated_features, gated_item], axis=1)
        outs = tf.matmul(outs, self.kernel['W_out'])
        if self.out_use_bias:
            outs += self.bias['b_out']
        if self.out_activation is not None:
            outs = self.out_activation(outs)
        return tf.expand_dims(outs, 0)

    def compute_output_shape(self, input_shape):
        assert input_shape and len(input_shape) == 3
        assert input_shape[-1]
        output_shapes = list(input_shape)
        output_shapes[-1] = self.units
        return tuple(output_shapes)

    def get_config(self):
        config = {'units': self.units,
                  'activation': activations.serialize(self.activation),
                  'out_activation': activations.serialize(self.out_activation),
                  'use_bias': self.use_bias,
                  'out_use_bias': self.out_use_bias,
                  'kernel_initializer': initializers.serialize(
                      self.kernel_initializer),
                  'bias_initializer': initializers.serialize(
                      self.bias_initializer),
                  'kernel_regularizer': regularizers.serialize(
                      self.kernel_regularizer),
                  'bias_regularizer': regularizers.serialize(
                      self.bias_regularizer),
                  'activity_regularizer': regularizers.serialize(
                      self.activity_regularizer),
                  'kernel_constraint': constraints.serialize(
                      self.kernel_constraint),
                  'bias_constraint': constraints.serialize(self.bias_constraint),
                  'features': self.features.numpy()
                  }

        base_config = super(CrossGate, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
